package jsonx

import (
	"encoding/json"
	"errors"
	"reflect"
	"sync"
)

var decoders = make(map[reflect.Type]map[any]reflect.Type)
var lock sync.RWMutex

type IJsonX[K comparable] interface {
	GetJsonXType() K
}

func RegisterJsonInterfaceType[K comparable, H IJsonX[K]](schema H) {
	lock.Lock()
	defer lock.Unlock()
	jsonXType := schema.GetJsonXType()
	idToType, ex := decoders[reflect.TypeOf(jsonXType)]
	if !ex {
		idToType = make(map[any]reflect.Type)
		decoders[reflect.TypeOf(jsonXType)] = idToType
	}
	idToType[jsonXType] = reflect.TypeOf(schema)
}

func getJsonInterfaceType(id any) reflect.Type {
	lock.RLock()
	defer lock.RUnlock()
	idToType, ex := decoders[reflect.TypeOf(id)]
	if !ex {
		return nil
	}
	return idToType[id]
}

type interfaceData[K comparable] struct {
	Id   K
	Data []byte
}

type JsonInterface[K comparable, T IJsonX[K]] struct {
	Value T `json:"-"`
}

func (c *JsonInterface[K, T]) MarshalJSON() ([]byte, error) {
	bs, err := json.Marshal(c.Value)
	if err != nil {
		panic(err)
	}
	rawData := interfaceData[K]{
		Id:   c.Value.GetJsonXType(),
		Data: bs,
	}
	return json.Marshal(rawData)
}

func (c *JsonInterface[K, T]) UnmarshalJSON(bytes []byte) error {
	rawData := &interfaceData[K]{}
	err := json.Unmarshal(bytes, rawData)
	if err != nil {
		return err
	}
	c.Value = *new(T)
	rt := getJsonInterfaceType(rawData.Id)
	if rt == nil {
		return errors.New("json interface not registered")
	}
	v := reflect.New(rt)
	err = json.Unmarshal(rawData.Data, v.Interface())
	if err != nil {
		return err
	}
	c.Value = v.Elem().Interface().(T)
	return nil
}
